package com.zj.esproxy.wrapper;

import cn.hutool.core.collection.CollectionUtil;
import com.alibaba.fastjson.JSONObject;
import com.alibaba.fastjson.annotation.JSONField;
import com.zj.baseproxy.utils.ClassUtil;
import com.zj.baseproxy.utils.LambdaUtil;
import lombok.Getter;
import org.apache.commons.lang3.StringUtils;
import org.apache.commons.lang3.reflect.FieldUtils;
import org.springframework.core.annotation.AnnotatedElementUtils;

import java.lang.reflect.Field;
import java.util.*;
import java.util.function.Consumer;

/**
 * @author arthur_zhou
 * @date 2023/11/7 14:08
 */
@Getter
public class EsWrapper {

    @JSONField(name = "query_id")
    private String queryId;

    private String project;

    private Integer timeout;

    private String product;

    private String operation;

    private Integer pos;

    private Integer count;

    private JSONObject filter;

    @JSONField(name = "display_field")
    private List<String> displayField;

    private List<AggWrapper> aggs;

    private Collapse collapse;

    private Key key;

    private List<Sort> sort;

    private Map<String, List<Highlight>> highlight;

    public EsWrapper(String project) {
        this.project = project;
    }

    public EsWrapper(EsWrapper esWrapper) {
        if (Objects.isNull(esWrapper)) {
            return;
        }
        if (StringUtils.isNotEmpty(esWrapper.project)) {
            this.project = esWrapper.project;
        }
        if (Objects.nonNull(esWrapper.timeout)) {
            this.timeout = esWrapper.timeout;
        }
        if (StringUtils.isNotEmpty(esWrapper.product)) {
            this.product = esWrapper.product;
        }
        if (StringUtils.isNotEmpty(esWrapper.operation)) {
            this.operation = esWrapper.operation;
        }
        if (Objects.nonNull(esWrapper.pos)) {
            this.pos = esWrapper.pos;
        }
        if (Objects.nonNull(esWrapper.count)) {
            this.count = esWrapper.count;
        }
        if (CollectionUtil.isNotEmpty(esWrapper.filter)) {
            this.filter = new JSONObject(new LinkedHashMap<>(esWrapper.filter));
        }
        if (CollectionUtil.isNotEmpty(esWrapper.displayField)) {
            this.displayField = new ArrayList<>(esWrapper.displayField);
        }
        if (CollectionUtil.isNotEmpty(esWrapper.aggs)) {
            this.aggs = new ArrayList<>(esWrapper.aggs);
        }
        if (Objects.nonNull(esWrapper.collapse)) {
            this.collapse = new Collapse(esWrapper.collapse);
        }
    }

    /**
     * 拼接filter条件
     *
     * @param wrapper 查询拼接器
     * @return 自己
     */
    public EsWrapper query(EsQueryWrapper wrapper) {
        if (Objects.nonNull(wrapper)) {
            filter = wrapper.build();
        }
        return this;
    }

    /**
     * 拼接filter条件
     *
     * @param consumer 函数
     * @return 自己
     */
    public EsWrapper query(Consumer<EsQueryWrapper> consumer) {
        EsQueryWrapper wrapper = new EsQueryWrapper();
        consumer.accept(wrapper);
        return query(wrapper);
    }

    /**
     * 搜索请求id
     * 不指定，默认用请求头里的request-id
     *
     * @param queryId 搜索请求id
     * @return 自己
     */
    public EsWrapper queryId(String queryId) {
        this.queryId = queryId;
        return this;
    }

    /**
     * 超时返回
     *
     * @param timeout 毫秒
     * @return 自己
     */
    public EsWrapper timeout(Integer timeout) {
        this.timeout = timeout;
        return this;
    }

    /**
     * 超时返回
     *
     * @param product 产品名称
     * @return 自己
     */
    public EsWrapper product(String product) {
        this.product = product;
        return this;
    }

    /**
     * 超时返回
     *
     * @param operation 操作名称
     * @return 自己
     */
    public EsWrapper operation(String operation) {
        this.operation = operation;
        return this;
    }

    /**
     * 起始下标
     *
     * @param pos 下标
     * @return 自己
     */
    public EsWrapper pos(Integer pos) {
        this.pos = pos;
        return this;
    }

    /**
     * 返回条数
     *
     * @param count 返回条数
     * @return 自己
     */
    public EsWrapper count(Integer count) {
        this.count = count;
        return this;
    }

    /**
     * 分页
     *
     * @param pos   下标
     * @param count 返回条数
     * @return 自己
     */
    public EsWrapper limit(Integer pos, Integer count) {
        this.pos = pos;
        this.count = count;
        return this;
    }

    /**
     * 聚合
     *
     * @param consumer 函数
     * @return 自己
     */
    public EsWrapper agg(Consumer<AggWrapper> consumer) {
        AggWrapper aggWrapper = new AggWrapper();
        consumer.accept(aggWrapper);
        return agg(aggWrapper);
    }

    /**
     * 聚合
     *
     * @param wrapper 拼接器
     * @return 自己
     */
    public EsWrapper agg(AggWrapper wrapper) {
        if (Objects.nonNull(wrapper) && Objects.nonNull(wrapper.getAggType())) {
            if (Objects.isNull(this.getAggs())) {
                this.aggs = new ArrayList<>();
            }
            this.getAggs().add(wrapper);
        }
        return this;
    }

    /**
     * 排序
     *
     * @param field 排序字段
     * @return 自己
     */
    public EsWrapper sort(String field) {
        return sort(field, null, null, null);
    }

    /**
     * 排序
     *
     * @param field 排序字段
     * @param desc  是否倒序
     * @return 自己
     */
    public EsWrapper sort(String field, Boolean desc) {
        return sort(field, desc, null, null);
    }

    /**
     * 排序
     *
     * @param field 排序字段
     * @param desc  是否倒序
     * @param lon   经度
     * @param lat   纬度
     * @return 自己
     */
    public EsWrapper sort(String field, Boolean desc, Object lon, Object lat) {
        return sort(new Sort(field, desc, lon, lat));
    }

    /**
     * 排序
     *
     * @param sort 排序信息
     * @return 自己
     */
    public EsWrapper sort(Sort sort) {
        if (Objects.nonNull(sort)) {
            if (Objects.isNull(this.sort)) {
                this.sort = new ArrayList<>();
            }
            this.sort.add(sort);
        }
        return this;
    }

    /**
     * 高亮
     *
     * @param field 高亮字段
     * @return 自己
     */
    public EsWrapper highlight(String field) {
        return highlight(field, null, null);
    }

    /**
     * 高亮
     *
     * @param field             高亮字段
     * @param fragmentSize      返回的高亮数据周围携带多少字符
     * @param numberOfFragments 返回多少个高亮块
     * @return 自己
     */
    public EsWrapper highlight(String field, Integer fragmentSize, Integer numberOfFragments) {
        return highlight(new Highlight(field, fragmentSize, numberOfFragments));
    }

    /**
     * 高亮
     *
     * @param highlight 高亮信息
     * @return 自己
     */
    public EsWrapper highlight(Highlight highlight) {
        if (Objects.nonNull(highlight)) {
            if (Objects.isNull(this.highlight)) {
                this.highlight = Collections.singletonMap("fields", new ArrayList<>());
            }
            this.highlight.get("fields").add(highlight);
        }
        return this;
    }

    /**
     * 高亮
     *
     * @param query           查询值
     * @param defaultOperator 分词关系
     * @param fields          查询字段
     * @return 自己
     */
    public EsWrapper key(String query, String defaultOperator, Key.Field... fields) {
        return key(new Key(query, defaultOperator, Objects.nonNull(fields) ? Arrays.asList(fields) : null));
    }

    /**
     * 高亮
     *
     * @param query           查询值
     * @param defaultOperator 分词关系
     * @param fields          查询字段
     * @return 自己
     */
    public EsWrapper key(String query, String defaultOperator, List<Key.Field> fields) {
        return key(new Key(query, defaultOperator, fields));
    }

    /**
     * 高亮
     *
     * @param key key信息
     * @return 自己
     */
    public EsWrapper key(Key key) {
        this.key = key;
        return this;
    }

    /**
     * 展示字段
     *
     * @param fields 字段数组
     * @return 自己
     */
    public EsWrapper displayField(String... fields) {
        return displayField(Objects.nonNull(fields) ? Arrays.asList(fields) : null);
    }

    /**
     * 展示字段
     *
     * @param fields 字段数组
     * @return 自己
     */
    public EsWrapper displayField(List<String> fields) {
        if (CollectionUtil.isNotEmpty(fields)) {
            if (Objects.isNull(this.displayField)) {
                this.displayField = new ArrayList<>();
            }
            this.displayField.addAll(fields);
        }
        return this;
    }

    /**
     * 展示字段
     *
     * @param clazz 实体类class
     * @return 自己
     */
    public EsWrapper displayField(Class<?> clazz) {
        if (Objects.nonNull(clazz) && !ClassUtil.isBaseType(clazz)) {
            Field[] allFields = FieldUtils.getAllFields(clazz);
            if (Objects.nonNull(allFields) && allFields.length > 0) {
                if (Objects.isNull(this.displayField)) {
                    this.displayField = new ArrayList<>();
                }
                for (Field field : allFields) {
                    JSONField jsonField = AnnotatedElementUtils.getMergedAnnotation(field, JSONField.class);
                    String filed;
                    if (Objects.isNull(jsonField) || StringUtils.isEmpty(filed = jsonField.name())) {
                        filed = LambdaUtil.camelToUnderscore(field.getName());
                    }
                    this.displayField.add(filed);
                }
            }
        }
        return this;
    }

    /**
     * 折叠
     *
     * @param field 折叠字段
     * @return 自己
     */
    public EsWrapper collapse(String field) {
        return collapse(field, (List<Collapse.InnerHit>) null);
    }


    /**
     * 折叠
     *
     * @param field     折叠字段
     * @param innerHits 折叠内展示
     * @return 自己
     */
    public EsWrapper collapse(String field, Collapse.InnerHit... innerHits) {
        return collapse(field, Objects.nonNull(innerHits) ? Arrays.asList(innerHits) : null);
    }

    /**
     * 折叠
     *
     * @param field     折叠字段
     * @param innerHits 折叠内展示
     * @return 自己
     */
    public EsWrapper collapse(String field, List<Collapse.InnerHit> innerHits) {
        return collapse(new Collapse(field, innerHits));
    }

    /**
     * 折叠
     *
     * @param collapse collapse信息
     * @return 自己
     */
    public EsWrapper collapse(Collapse collapse) {
        this.collapse = collapse;
        return this;
    }

}
